use std::collections::HashMap;
use std::fs::File;
use std::io::{Read, Result};
use std::path::{Path, PathBuf};

use ndarray::Array2;
use ndarray_npy::NpzReader;
use regex::Regex;

fn scan_folder(base_path: &Path) -> Result<Vec<String>> {
    let number_re = Regex::new(r"^\d+$").unwrap();
    let output = std::fs::read_dir(base_path)?;
    let mut ret = Vec::new();
    for entry in output {
        let entry = entry?;
        let file_type = entry.file_type()?;
        if file_type.is_dir() {
            let filename = entry.file_name().to_string_lossy().to_string();
            if number_re.is_match(&filename) {
                ret.push(filename);
            }
        }
    }
    Ok(ret)
}

fn scan_frames(base_path: &Path) -> Result<Vec<String>> {
    let number_re = Regex::new(r"^\d+\.jpg$").unwrap();
    let output = std::fs::read_dir(base_path)?;
    let mut ret = Vec::new();
    for entry in output {
        let entry = entry?;
        let file_type = entry.file_type()?;
        if file_type.is_file() {
            let filename = entry.file_name().to_string_lossy().to_string();
            if number_re.is_match(&filename) {
                ret.push(filename);
            }
        }
    }
    ret.sort_by_key(|e| {
        let id = u32::from_str_radix(&e[0..e.len() - 4], 10).expect("invalid filename");
        id
    });

    let ret = ret
        .into_iter()
        .map(|f| {
            let mut full = PathBuf::from(base_path);
            full.push(&f);
            full.to_string_lossy().to_string()
        })
        .collect::<Vec<_>>();
    Ok(ret)
}

pub fn scan_dataset(base_path: &Path) -> Result<HashMap<u32, Vec<String>>> {
    let folders = scan_folder(base_path)?;
    let mut ret = HashMap::new();

    for f in folders {
        let video_id = u32::from_str_radix(&f, 10).expect("invalid video id");
        let mut video_base_path = PathBuf::from(base_path);
        video_base_path.push(f);
        // println!("scan base path in {:?}", video_base_path);
        let frames = scan_frames(&video_base_path)?;
        ret.insert(video_id, frames);
    }
    Ok(ret)
}

pub struct Dataset {
    pub videos: HashMap<u32, Vec<String>>,
    pub label: LabelStore,
    pub infer: Option<InferStore>,
    base_path: PathBuf,
}

impl Dataset {

    pub fn load_label(&self) -> String {
        let mut label_path = self.base_path.clone();
        label_path.push("keyframe.json");
        if label_path.is_file() {
            std::fs::read_to_string(&label_path).expect("read label path failed")
        } else {
            "{}".to_owned()
        }
    }

    pub fn save_label(&self, content: String) {
        let mut label_path = self.base_path.clone();
        label_path.push("keyframe.json");
        std::fs::write(label_path, content).expect("write label failed");
    }

}

pub fn parse_dataset(base_path: &str, infer_path: Option<&str>) -> Result<Dataset> {
    let mut path = PathBuf::from(base_path);
    let base_path = path.clone();
    let mut videos = scan_dataset(&path)?;

    let mut target_path = path.clone();
    target_path.push("target.json");

    if target_path.is_file() {
        let target_json_content = std::fs::read_to_string(&target_path).unwrap();
        let video_id_list: Vec<u32> = serde_json::from_str(&target_json_content).expect("invalid target.json");

        let invalid_id: Vec<u32> = videos.keys().filter(|k| !video_id_list.contains(*k)).cloned().collect();
        for id in invalid_id {
            videos.remove(&id);
        }

    }

    path.push("labels.npz");
    let label = load_label(&path);

    let infer = if let Some(path) = infer_path {
        let path = PathBuf::from(path);
        let infer = InferStore::load(&path);
        Some(infer)
    } else {
        None
    };

    Ok(Dataset { videos, label, infer, base_path })
}

pub struct LabelStore {
    offset: HashMap<u64, u64>,
    label: Array2<f32>,
    raw: Array2<f32>,
}

impl LabelStore {
    pub fn fetch_raw(&self, video_id: u64, frame_id: u64) -> [f32; 4] {
        let v_offset = self.offset[&video_id];
        let offset = (v_offset + frame_id) as usize;

        let mut ret = [0f32; 4];
        ret[0] = *self.raw.get([offset, 0]).unwrap();
        ret[1] = *self.raw.get([offset, 1]).unwrap();
        ret[2] = *self.raw.get([offset, 2]).unwrap();
        ret[3] = *self.raw.get([offset, 3]).unwrap();

        ret
    }
}

pub fn load_label(base_path: &Path) -> LabelStore {
    println!("load label from: {:?}", base_path);

    let mut npz = NpzReader::new(File::open(base_path).unwrap()).unwrap();

    let names = npz.names().unwrap();
    for name in names {
        println!("found name: {}", name);
    }

    let offset_array: Array2<i64> = npz.by_name("offset.npy").unwrap();
    let raw: Array2<f32> = npz.by_name("raw_quat.npy").unwrap();
    let label: Array2<f32> = npz.by_name("label.npy").unwrap();

    println!("offset array shape: {:?}", offset_array.shape());
    println!("raw array shape: {:?}", raw.shape());
    println!("label array shape: {:?}", label.shape());

    let mut offset = HashMap::new();
    for i in 0..offset_array.shape()[0] {
        let video_id = *offset_array.get([i, 0]).unwrap();
        let video_offset = *offset_array.get([i, 1]).unwrap();
        offset.insert(video_id as u64, video_offset as u64);
    }

    LabelStore { offset, label, raw }
}

pub struct InferStore {
    index_map: HashMap<(u64, u64), usize>,
    preds_quats: Array2<f32>,
}

impl InferStore {
    pub fn load(path: &Path) -> InferStore {
        println!("load infer from: {:?}", path);

        let mut npz = NpzReader::new(File::open(path).unwrap()).unwrap();

        let names = npz.names().unwrap();
        for name in names {
            println!("found name: {}", name);
        }

        let indexes: Array2<i64> = npz.by_name("indexes.npy").unwrap();
        let preds_quats: Array2<f32> = npz.by_name("preds_quat.npy").unwrap();

        println!("indexes array shape: {:?}", indexes.shape());
        println!("preds_quats array shape: {:?}", preds_quats.shape());

        let mut index_map = HashMap::new();
        for i in 0..indexes.shape()[0] {
            let preds_idx = i;
            let video_id = *indexes.get([i, 0]).unwrap();
            let frame_id = *indexes.get([i, 1]).unwrap();
            index_map.insert((video_id as u64, frame_id as u64), preds_idx);
        }

        InferStore { index_map, preds_quats }
    }

    pub fn find_pred(&self, video_id: u64, frame_id: u64) -> Option<[f32; 4]> {
        if let Some(index) = self.index_map.get(&(video_id, frame_id)) {
            let index = *index;
            let mut ret = [0f32; 4];
            ret[0] = *self.preds_quats.get([index, 0]).unwrap();
            ret[1] = *self.preds_quats.get([index, 1]).unwrap();
            ret[2] = *self.preds_quats.get([index, 2]).unwrap();
            ret[3] = *self.preds_quats.get([index, 3]).unwrap();
            Some(ret)
        } else {
            None
        }
    }
}
